/*
 * Copyright (c) 2025 [ByteDance Ltd. and/or its affiliates.] and/or its affiliates
 * SPDX-License-Identifier: LGPL v2.1+
*/
#include <stdio.h>
#include <string.h>

#include "libavutil/opt.h"
#include "libavutil/eval.h"

#include "avfilter.h"
#include "formats.h"
#include "internal.h"
#include "video.h"

#import <MetalPerformanceShaders/MetalPerformanceShaders.h>
#import <CoreVideo/CoreVideo.h>
#import <CoreVideo/CVPixelBuffer.h>
#import <Metal/MTLDevice.h>
#import <Metal/MTLBuffer.h>
#import <Metal/MTLTexture.h>

typedef struct VTScaleContext {
    const AVClass *class;
    
    id<MTLDevice> mtlDevice;
    id<MTLCommandQueue> mtlCommandQueue;
    
    MPSImageLanczosScale *lanczosScale;
    CVPixelBufferRef outPixelbuf;
    /**
     * New dimensions. Special values are:
     *   0 = original width/height
     *  -1 = keep original aspect
     */
    int w, h;
    
    /**
     * Output sw format. AV_PIX_FMT_NONE for no conversion.
     */
    enum AVPixelFormat format;

    char *w_expr;
    char *h_expr;
    char *scale_type_expr;
} VTScaleContext;

static const char *const var_names[] = {
    "out_w",  "ow",
    "out_h",  "oh",
    NULL
};

enum var_name {
    VAR_OUT_W,  VAR_OW,
    VAR_OUT_H,  VAR_OH,
    VARS_NB
};

static av_cold int metalscale_init(AVFilterContext *ctx) {
    VTScaleContext* vtCtx = ctx->priv;
    vtCtx->mtlDevice = MTLCreateSystemDefaultDevice();
    vtCtx->mtlCommandQueue = [vtCtx->mtlDevice newCommandQueue];
    
    if (!strcmp(vtCtx->scale_type_expr, "lanczos")) {
        vtCtx->lanczosScale = [[MPSImageLanczosScale alloc] initWithDevice:vtCtx->mtlDevice];
    } else {
        av_log(ctx, AV_LOG_ERROR, "unsupported scale mode");
        return AVERROR(ESRCH);
    }
    
    return 0;
}

static av_cold int metalscale_uninit(AVFilterContext *ctx) {
    VTScaleContext* vtCtx = ctx->priv;
    if (!vtCtx) {
        return -1;
    }
    
    if (vtCtx->lanczosScale) {
        [vtCtx->lanczosScale release];
        vtCtx->lanczosScale = nil;
    }
    
    if (vtCtx->outPixelbuf) {
        CVPixelBufferRelease(vtCtx->outPixelbuf);
        vtCtx->outPixelbuf = nil;
    }
    
    if (vtCtx->mtlCommandQueue) {
        [vtCtx->mtlCommandQueue release];
        vtCtx->mtlCommandQueue = nil;
    }
    
    if (vtCtx->mtlDevice) {
        [vtCtx->mtlDevice release];
        vtCtx->mtlDevice = nil;
    }
    return 0;
}

#define OFFSET(x) offsetof(VTScaleContext, x)
#define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
static const AVOption options[] = {
    { "w",      "Output video width",  OFFSET(w_expr),     AV_OPT_TYPE_STRING, { .str = "iw"   }, .flags = FLAGS },
    { "h",      "Output video height", OFFSET(h_expr),     AV_OPT_TYPE_STRING, { .str = "ih"   }, .flags = FLAGS },
    { "scale_type", "Output scale type", OFFSET(scale_type_expr), AV_OPT_TYPE_STRING, { .str = "lanczos" }, .flags = FLAGS },
    { NULL },
};

static const AVClass metalscale_class = {
    .class_name = "scale_metal",
    .item_name  = av_default_item_name,
    .category   = AV_CLASS_CATEGORY_FILTER,
    .option     = options,
    .version    = LIBAVUTIL_VERSION_INT,
};

static int metalsacle_getset_colorinfo(CVPixelBufferRef inPixelbuf, CVPixelBufferRef outPixelbuf) {
    @autoreleasepool {
        if (!inPixelbuf || !outPixelbuf) {
            return -1;
        }

        CFTypeRef color_matrix = CVBufferGetAttachment(inPixelbuf, kCVImageBufferYCbCrMatrixKey, NULL);
        CFTypeRef color_primary = CVBufferGetAttachment(inPixelbuf, kCVImageBufferColorPrimariesKey, NULL);
        CFTypeRef color_transfer = CVBufferGetAttachment(inPixelbuf, kCVImageBufferTransferFunctionKey, NULL);

        CVBufferSetAttachment(outPixelbuf, kCVImageBufferYCbCrMatrixKey, (CFStringRef)color_matrix, kCVAttachmentMode_ShouldPropagate);
        CVBufferSetAttachment(outPixelbuf, kCVImageBufferColorPrimariesKey, (CFStringRef)color_primary, kCVAttachmentMode_ShouldPropagate);
        CVBufferSetAttachment(outPixelbuf, kCVImageBufferTransferFunctionKey, (CFStringRef)color_transfer, kCVAttachmentMode_ShouldPropagate);
    }

    return 0;
}

static int metalscale_config_props(AVFilterLink *outlink) {
    AVFilterContext *ctx = outlink->src;
    AVFilterLink *inlink = outlink->src->inputs[0];
    VTScaleContext *s = ctx->priv;
    int64_t w, h;
    double var_values[VARS_NB], res;
    char *expr;
    int ret;

    av_expr_parse_and_eval(&res, (expr = s->w_expr),
                          var_names, var_values,
                          NULL, NULL, NULL, NULL, NULL, 0, ctx);
    s->w = var_values[VAR_OUT_W] = var_values[VAR_OW] = res;
    if ((ret = av_expr_parse_and_eval(&res, (expr = s->h_expr),
                                     var_names, var_values,
                                     NULL, NULL, NULL, NULL, NULL, 0, ctx)) < 0)
       goto fail;
    s->h = var_values[VAR_OUT_H] = var_values[VAR_OH] = res;

    if ((ret = av_expr_parse_and_eval(&res, (expr = s->w_expr),
                                     var_names, var_values,
                                     NULL, NULL, NULL, NULL, NULL, 0, ctx)) < 0)
       goto fail;
    s->w = res;

    w = s->w;
    h = s->h;

    if (w <  -1 || h <  -1) {
       av_log(ctx, AV_LOG_ERROR, "error, invalid metal scale size\n");
       return AVERROR(EINVAL);
    }
    if (w == -1 && h == -1)
       s->w = s->h = 0;

    if (!(w = s->w))
       w = inlink->w;
    if (!(h = s->h))
       h = inlink->h;
    if (w == -1)
       w = av_rescale(h, inlink->w, inlink->h);
    if (h == -1)
       h = av_rescale(w, inlink->h, inlink->w);

    outlink->w = w;
    outlink->h = h;

    av_log(ctx, AV_LOG_VERBOSE, "w:%d h:%d -> w:%d h:%d\n",
          inlink->w, inlink->h, outlink->w, outlink->h);

    return 0;

fail:
    av_log(ctx, AV_LOG_ERROR,
          "error when evaluating the expression '%s'\n", expr);
    return ret;


    return 0;
}

static int metalscale_filter_frame(AVFilterLink* link, AVFrame* in) {
    AVFilterContext* ctx = link->dst;
    VTScaleContext* vtCtx = ctx->priv;
    AVFilterLink* outlink = ctx->outputs[0];
    
    AVFrame *out = NULL;
    out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
    if (!out) {
        out = av_frame_alloc();
        out->width   = vtCtx->w;
        out->height  = vtCtx->h;
        out->format  = vtCtx->format;
        av_frame_get_buffer(out, 0);
    }
 @autoreleasepool {
        CVPixelBufferRef inImg = (CVPixelBufferRef)in->data[3];
        if (!inImg) {
            return AVERROR(EIO);
        }

        int16_t width = CVPixelBufferGetWidth(inImg);
        int16_t height = CVPixelBufferGetHeight(inImg);
        OSType pixelFormat = CVPixelBufferGetPixelFormatType(inImg);
        bool srcIs10bit = false;
        if (pixelFormat == kCVPixelFormatType_64RGBALE) {
            srcIs10bit = true;
        }
        MTLPixelFormat textureFormat = srcIs10bit ? MTLPixelFormatRGBA16Unorm : MTLPixelFormatBGRA8Unorm;

        CVMetalTextureRef inTexture = NULL;
        CVMetalTextureCacheRef textureCache = NULL;
        CVReturn crtCacheStatus = CVMetalTextureCacheCreate(kCFAllocatorDefault, NULL, vtCtx->mtlDevice, NULL, &textureCache);
        if (crtCacheStatus != kCVReturnSuccess) {
            return AVERROR(ENOMEM);
        }

        CVReturn inStatus = CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache, inImg, NULL, textureFormat, width, height, 0, &inTexture);
        if(inStatus != kCVReturnSuccess) {
            CVBufferRelease(inTexture);
            return AVERROR(ENOMEM);
        }

        CVMetalTextureRef outTexture = NULL;
        NSDictionary *options = @{(id)kCVPixelBufferIOSurfacePropertiesKey: @ {}};
        if (!vtCtx->outPixelbuf) {
            CVReturn crtStatus = CVPixelBufferCreate(kCFAllocatorDefault, outlink->w, outlink->h, pixelFormat, (__bridge CFDictionaryRef)(options), &vtCtx->outPixelbuf);
            if (crtStatus != kCVReturnSuccess) {
               CVBufferRelease(inTexture);
               return AVERROR(ENOMEM);
            }
        }
        

        int colorStatus = metalsacle_getset_colorinfo(inImg, vtCtx->outPixelbuf);
        if (colorStatus == -1) {
            CVBufferRelease(inTexture);
            av_log(ctx, AV_LOG_ERROR,
                "get set color info error");
                return AVERROR(EIO);
        }

        CVReturn outStatus = CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, textureCache, vtCtx->outPixelbuf, NULL, textureFormat, outlink->w, outlink->h, 0, &outTexture);
        if(outStatus != kCVReturnSuccess) {
            CVBufferRelease(inTexture);
            CVBufferRelease(outTexture);
            return AVERROR(ENOMEM);
        }

        id<MTLCommandBuffer> commandBuffer = [vtCtx->mtlCommandQueue commandBuffer];
        id<MTLTexture> mtlInTexture = CVMetalTextureGetTexture(inTexture);
        id<MTLTexture> mtlOutTexture = CVMetalTextureGetTexture(outTexture);
        if (vtCtx->lanczosScale) {
            [vtCtx->lanczosScale encodeToCommandBuffer:commandBuffer sourceTexture:mtlInTexture destinationTexture:mtlOutTexture];
            [commandBuffer commit];
            [commandBuffer waitUntilCompleted];
        } else {
            av_log(ctx, AV_LOG_ERROR, "unsupported scale mode");
            CVBufferRelease(inTexture);
            CVBufferRelease(outTexture);
            return AVERROR(ESRCH);
        }

        CVBufferRelease(inTexture);
        CVBufferRelease(outTexture);

        out->data[3] = (uint8_t *)CVPixelBufferRetain(vtCtx->outPixelbuf);
        out->format  = AV_PIX_FMT_VIDEOTOOLBOX;
        if (av_frame_copy_props(out, in) < 0) {
            return AVERROR(ENOMEM);
        }

        if (!out)
            return AVERROR(ENOMEM);

        av_frame_free(&in);
    }
    return ff_filter_frame(outlink, out);
}

static const AVFilterPad metalscale_inputs[] = {
    {
        .name         = "default",
        .type         = AVMEDIA_TYPE_VIDEO,
        .filter_frame = metalscale_filter_frame,
    },
};

static const AVFilterPad metalscale_outputs[] = {
    {
        .name         = "default",
        .type         = AVMEDIA_TYPE_VIDEO,
        .config_props = metalscale_config_props,
    },
};

const AVFilter ff_vf_scale_metal = {
    .name = "scale_metal",
    .description = NULL_IF_CONFIG_SMALL("metal video scaling and format conversion"),
    .priv_size = sizeof(VTScaleContext),
    .init = metalscale_init,
    .uninit = metalscale_uninit,
    FILTER_INPUTS(metalscale_inputs),
    FILTER_OUTPUTS(metalscale_outputs),
    FILTER_SINGLE_PIXFMT(AV_PIX_FMT_VIDEOTOOLBOX),
    .priv_class = &metalscale_class,
};
